from .base_exchange import *
import gate_api
from gate_api import SpotApi, FuturesApi, Configuration
import asyncio
import websockets
import json

class GateExchange(BaseExchange):
    def __init__(self, api_key: str = "", api_secret: str = "", is_futures: bool = False):
        super().__init__(api_key, api_secret)
        self.name = "Gate.io"
        self.is_futures = is_futures
        
        # 初始化API
        self.config = Configuration(
            key=api_key,
            secret=api_secret,
            host="https://api.gateio.ws/api/v4"
        )
        self.spot_api = SpotApi(gate_api.ApiClient(self.config))
        self.futures_api = FuturesApi(gate_api.ApiClient(self.config))
        
        # WebSocket连接
        self.ws = None
        self.ws_connected = False
        self.subscriptions = set()
        
    async def connect_ws(self):
        """连接WebSocket"""
        if self.is_futures:
            ws_url = "wss://fx-ws.gateio.ws/v4/ws/usdt"
        else:
            ws_url = "wss://api.gateio.ws/ws/v4/"
            
        self.ws = await websockets.connect(ws_url)
        self.ws_connected = True
        
        # 启动消息处理
        asyncio.create_task(self._handle_ws_messages())
        
    async def _handle_ws_messages(self):
        """处理WebSocket消息"""
        while True:
            try:
                message = await self.ws.recv()
                data = json.loads(message)
                await self._process_ws_message(data)
            except Exception as e:
                print(f"WebSocket error: {e}")
                self.ws_connected = False
                await self.connect_ws()
                
    async def _process_ws_message(self, data: dict):
        """处理不同类型的WebSocket消息"""
        if 'channel' not in data:
            return
            
        channel = data['channel']
        if channel == 'spot.tickers' or channel == 'futures.tickers':
            await self._handle_ticker_update(data)
        elif channel == 'spot.orders' or channel == 'futures.orders':
            await self._handle_order_update(data)
            
    async def get_ticker(self, symbol: str) -> Ticker:
        """获取最新行情"""
        try:
            if self.is_futures:
                response = await self.futures_api.list_futures_tickers(settle='usdt', contract=symbol)
                return Ticker(
                    symbol=symbol,
                    last_price=float(response.last),
                    bid_price=float(response.bid),
                    ask_price=float(response.ask),
                    volume_24h=float(response.volume_24h),
                    timestamp=time.time()
                )
            else:
                response = await self.spot_api.list_tickers(currency_pair=symbol)
                return Ticker(
                    symbol=symbol,
                    last_price=float(response[0].last),
                    bid_price=float(response[0].highest_bid),
                    ask_price=float(response[0].lowest_ask),
                    volume_24h=float(response[0].base_volume),
                    timestamp=time.time()
                )
        except Exception as e:
            raise Exception(f"Failed to get ticker: {str(e)}") 